# -*- coding: utf-8 -*-
"""
处理射击的相关程序
"""

__authur__ = 'Yu XH'

import time

from myAI.attdamage.weapon_info import *
from myAI.const import *
from myAI.hexmap.handle_map import *
from myAI.attdamage.ruling_data import *


class AttLevelRect:
    # ElevDiff为100*200的矩阵，行为高差（绝对值），列为距离，都从0开始算
    ElevDiff = np.concatenate(
        (np.concatenate(
            (np.array([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
                       [0, -2, -2, -1, -1, -1, 0, 0, 0, 0, 0, 0, 0],
                       [0, -2, -2, -2, -1, -1, -1, -1, 0, 0, 0, 0, 0],
                       [0, -3, -2, -2, -2, -1, -1, -1, -1, -1, 0, 0, 0],
                       [0, -3, -3, -3, -2, -2, -1, -1, -1, -1, -1, -1, 0],
                       [0, -4, -3, -3, -3, -2, -2, -2, -1, -1, -1, -1, -1],
                       [0, -4, -4, -4, -3, -3, -2, -2, -2, -1, -1, -1, -1],
                       [0, -5, -4, -4, -4, -3, -2, -2, -2, -2, -1, -1, -1],
                       [0, -5, -5, -5, -4, -3, -3, -2, -2, -2, -2, -2, -1]]),
             np.tile(np.array([0, -5, -5, -5, -4, -3, -3, -2, -2, -2, -2, -2, -1]).reshape(1, -1), (91, 1))),
            axis=0),
         np.zeros((100, 187), dtype=int)),
        axis=1)


class TerrainMeanRect:
    """不同地形对损伤结果的期望修正影响"""
    # 地形顺序：OpenLand, Forest, Resident, SoftLand, Water
    Soldier = [0.42, 0.14, 0.28, 0.42, 0.42]
    Vehicle = {
        ArmorType.NoArmor: [1.17, 0.86, 1, 1.17, 0.86],
        ArmorType.Light: [0.92, 0.44, 0.67, 0.92, 0.44],
        ArmorType.Medium: [0.39, 0.08, 0.19, 0.39, 0.08],
        ArmorType.Heavy: [0.25, -0.06, 0.08, 0.25, -0.06],
        ArmorType.Composit: [-0.19, -0.78, -0.47, -0.19, -0.78]
    }


def get_weapon(wp_id):
    """根据武器ID调取武器信息"""
    if wp_id == 36:
        return Biggun()
    elif wp_id == 37:
        return Middlegun()
    elif wp_id == 54:
        return Littlegun()
    elif wp_id == 56:
        return QuickFiringGun()
    elif wp_id == 35:
        return RocketLauncher()
    elif wp_id == 29:
        return SoldierWeapon()
    elif wp_id == 43:
        return LightWeapons()
    elif wp_id == 71:
        return PortableMissile()
    elif wp_id == 69:
        return VehicularMissile()
    elif wp_id == 84:
        return GunLaunchedMissile()
    elif wp_id == 73:
        return HeavyMissile()
    elif wp_id == 83 or wp_id == 74:
        return MediumMissile()
    elif wp_id == 75:
        return SmallMissile()
    elif wp_id == 76:
        return CruiseMissile()


def confirm_bullet_types(special_shoot=None):
    """根据射击任务的特殊性，确定可以使用的武器/弹药类型bullet_types，用于integrate_att_level函数"""
    if not special_shoot:
        return None  # 表示所有类型的武器都可以使用
    elif special_shoot == "move_shoot":
        return [BulletType.Convention]  # 行进间射击行为中只能使用常规弹药，不能使用导弹


def get_shoot_radius(carry_weapon_ids, remain_bullet_nums, tar_type, special_shoot=None):
    """计算有效射击半径（攻击等级不等于0）"""
    if not isinstance(carry_weapon_ids, list):
        carry_weapon_ids = [carry_weapon_ids]
    bullet_types = confirm_bullet_types(special_shoot)
    if not bullet_types:
        bullet_types = [BulletType.Convention, BulletType.HeavyMissile,
                        BulletType.MediumMissile, BulletType.LightMissile]
    if not isinstance(bullet_types, list):
        bullet_types = [bullet_types]
    dist = 0
    for wp_id in carry_weapon_ids:
        weapon = get_weapon(wp_id)
        if weapon.type in bullet_types:  # 符合指定的武器类型
            if remain_bullet_nums[weapon.type] > 0:  # 弹药有剩余
                if dist < weapon.shoot_range[tar_type]:
                    dist = weapon.shoot_range[tar_type]
    return dist


def integrate_att_level(carry_weapon_ids, remain_bullet_nums, tar_type, bullet_types=None):
    """ 2020-5-21
    整合多个武器的攻击等级矩阵
    :param carry_weapon_ids: {list}武器的ID列表，参见算子属性operator['carry_weapon_ids']
    :param remain_bullet_nums: {dist}弹药剩余量，参见算子属性operator['remain_bullet_nums']
    :param tar_type: {int}目标类型（步兵、车辆、空中算子），参见BopType, const.py
    :param bullet_types: 指定可用的武器类型, None表示所有类型都可以使用
    :return: 整合后的攻击等级矩阵
    """
    if not isinstance(carry_weapon_ids, list):
        carry_weapon_ids = [carry_weapon_ids]
    if not bullet_types:
        bullet_types = [BulletType.Convention, BulletType.HeavyMissile,
                        BulletType.MediumMissile, BulletType.LightMissile]
    if not isinstance(bullet_types, list):
        bullet_types = [bullet_types]
    if tar_type == BopType.Infantry:
        integrated_att_level = np.zeros((4, 21))
    else:  # 车辆 or 空中算子
        integrated_att_level = np.zeros(21)
    used_wp_ids = []
    for wp_id in carry_weapon_ids:
        weapon = get_weapon(wp_id)
        if weapon.shoot_range[tar_type] > 0:  # 武器使用于目标类型（步兵、车辆、空中算子）
            if weapon.type in bullet_types:  # 符合指定的武器类型
                if remain_bullet_nums[weapon.type] > 0:  # 弹药有剩余
                    used_wp_ids.append(wp_id)
                    if tar_type == BopType.Infantry:
                        integrated_att_level = np.maximum(integrated_att_level, weapon.att_level_soldier)
                    elif tar_type == BopType.Vehicle:
                        integrated_att_level = np.maximum(integrated_att_level, weapon.att_level_car)
                    elif tar_type == BopType.Aircraft:
                        integrated_att_level = np.maximum(integrated_att_level, weapon.att_level_aircraft)
    if used_wp_ids:
        return integrated_att_level
    else:
        # print("没有针对目标算子的可用武器！")
        return integrated_att_level  # 此时输出的是全零向量/矩阵


def calc_thread_att_level(myAgent, att_weapon_ids, att_bullet_nums, att_locs, tar_type, tar_locs=np.array([]),
                          att_obj_type=BopType.Vehicle, att_blood=3, special_shoot=None):
    """
    计算攻击算子att_obj在设定的攻击位置att_locs上对类型为tar_type的算子的威胁性，以攻击等级作为威胁性的度量
    :param myAgent:
    :param att_weapon_ids: {list} 武器列表
    :param att_bullet_nums: {dict} 剩余弹药量
    :param att_locs: {narray} 攻击位置，mat_ind的一维数组
    :param tar_type: {int}目标类型（步兵、车辆、空中算子），参见BopType, const.py
    :param tar_locs: {narray} 目标位置，mat_ind的一维数组
    :param att_blood: {int}攻击算子的车/班数，缺省为3
    :param special_shoot: 是否是特殊射击行为，目前只有一种特殊射击行为："move_shoot"
    :return: 输出att_obj在att_locs的各点对地图上所有六角格的att_level，六角格用mat_ind索引
    """
    map_size = myAgent.get_map_size()
    if isinstance(att_locs, int):
        att_locs = np.array([att_locs])
    if isinstance(tar_locs, int):
        tar_locs = np.array([tar_locs])
    if tar_locs.size == 0:
        n_hex = map_size['row'] * map_size['col']
        tar_locs = np.arange(n_hex)  # 地图上的所有六角格
    bullet_types = confirm_bullet_types(special_shoot)
    dist_matrix = myAgent.my_map['dist_matrix'][att_locs.reshape(-1, 1), tar_locs.reshape(1, -1)]
    sup_dist = map_size['row'] + map_size['col']  # 距离的上限
    extended_att_level = np.zeros(sup_dist)  # 保证所有距离上都可以索引
    if tar_type == BopType.Infantry:
        integrated_att_level = integrate_att_level(att_weapon_ids,
                                                   att_bullet_nums,
                                                   BopType.Infantry,
                                                   bullet_types)[att_blood - 1]
    else:
        integrated_att_level = integrate_att_level(att_weapon_ids,
                                                   att_bullet_nums,
                                                   tar_type,
                                                   bullet_types)
    extended_att_level[0:21] = integrated_att_level
    att_level_results = extended_att_level[dist_matrix]
    # 进行通视修正
    if att_obj_type == BopType.Aircraft and tar_type == BopType.Aircraft:
        see_matrix = myAgent.my_map['see_matrix'][1][att_locs.reshape(-1, 1), tar_locs.reshape(1, -1)]
    elif att_obj_type != BopType.Aircraft and tar_type != BopType.Aircraft:
        see_matrix = myAgent.my_map['see_matrix'][0][att_locs.reshape(-1, 1), tar_locs.reshape(1, -1)]
    elif att_obj_type == BopType.Aircraft and tar_type != BopType.Aircraft:
        see_matrix = myAgent.my_map['see_matrix'][2][att_locs.reshape(-1, 1), tar_locs.reshape(1, -1)]
    else:
        see_matrix = myAgent.my_map['see_matrix'][2].T[att_locs.reshape(-1, 1), tar_locs.reshape(1, -1)]
    att_level_results[np.logical_not(see_matrix)] = 0
    # 计算高差对攻击等级的修正
    if tar_type != BopType.Aircraft:
        elev_diff_matrix = myAgent.my_map['elev_diff_matrix'][att_locs.reshape(-1, 1), tar_locs.reshape(1, -1)].astype(
            int)
        if att_obj_type == BopType.Aircraft:
            elev_diff_matrix += 200  # 对空中算子的高程进行修正
        att_level_rect = AttLevelRect.ElevDiff[np.minimum(elev_diff_matrix, 0) // -20, dist_matrix]
        # 修正初始结果
        att_level_results += att_level_rect
        att_level_results = np.maximum(att_level_results, 0).astype(int)  # 攻击等级不可能为负值
    if att_level_results.size == 1:
        return att_level_results[0]
    return att_level_results


def calc_threat_levels(myAgent, att_weapon_ids, att_bullet_nums, att_loc_probs, tar_type, tar_locs=np.array([]),
                       att_obj_type=BopType.Vehicle, att_blood=3, special_shoot=None):
    """
    计算攻击算子att_obj在设定的攻击位置（位置概率为att_loc_probs）上对类型为tar_type的算子的威胁性，以攻击等级作为威胁性的度量
    :param myAgent:
    :param att_weapon_ids: {list} 武器列表
    :param att_bullet_nums: {dict} 剩余弹药量
    :param att_loc_probs: {narray} 攻击算子所处位置的概率，mat_ind的一维数组，数组维度为n_hexes
    :param tar_type: {int}目标类型（步兵、车辆、空中算子），参见BopType, const.py
    :param tar_locs: {narray} 目标位置，mat_ind的一维数组
    :param att_blood: {int}攻击算子的车/班数，缺省为3
    :param special_shoot: 是否是特殊射击行为，目前只有一种特殊射击行为："move_shoot"
    :return: 输出对地图上六角格tar_locs的威胁程度（与tar_locs同维度），六角格用mat_ind索引
    """
    map_size = myAgent.get_map_size()
    n_hexes = map_size['row'] * map_size['col']
    if isinstance(tar_locs, int):
        tar_locs = np.array([tar_locs])
    if tar_locs.size == 0:
        tar_locs = np.arange(n_hexes)
    bullet_types = confirm_bullet_types(special_shoot)
    att_locs = np.arange(n_hexes)[att_loc_probs > 0]
    dist_matrix = myAgent.my_map['dist_matrix'][att_locs.reshape(-1, 1), tar_locs.reshape(1, -1)]
    sup_dist = map_size['row'] + map_size['col']  # 距离的上限
    extended_att_level = np.zeros(sup_dist)  # 保证所有距离上都可以索引
    if tar_type == BopType.Infantry:
        integrated_att_level = integrate_att_level(att_weapon_ids,
                                                   att_bullet_nums,
                                                   BopType.Infantry,
                                                   bullet_types)[att_blood - 1]
    else:
        integrated_att_level = integrate_att_level(att_weapon_ids,
                                                   att_bullet_nums,
                                                   tar_type,
                                                   bullet_types)
    extended_att_level[0:21] = integrated_att_level
    att_level_results = extended_att_level[dist_matrix]
    # 进行通视修正
    if att_obj_type == BopType.Aircraft and tar_type == BopType.Aircraft:
        see_matrix = myAgent.my_map['see_matrix'][1][att_locs.reshape(-1, 1), tar_locs.reshape(1, -1)]
    elif att_obj_type != BopType.Aircraft and tar_type != BopType.Aircraft:
        see_matrix = myAgent.my_map['see_matrix'][0][att_locs.reshape(-1, 1), tar_locs.reshape(1, -1)]
    elif att_obj_type == BopType.Aircraft and tar_type != BopType.Aircraft:
        see_matrix = myAgent.my_map['see_matrix'][2][att_locs.reshape(-1, 1), tar_locs.reshape(1, -1)]
    else:
        see_matrix = myAgent.my_map['see_matrix'][2].T[att_locs.reshape(-1, 1), tar_locs.reshape(1, -1)]
    att_level_results[np.logical_not(see_matrix)] = 0
    # 计算高差对攻击等级的修正
    if tar_type != BopType.Aircraft:
        elev_diff_matrix = myAgent.my_map['elev_diff_matrix'][att_locs.reshape(-1, 1), tar_locs.reshape(1, -1)].astype(
            int)
        if att_obj_type == BopType.Aircraft:
            elev_diff_matrix += 200  # 对空中算子的高程进行修正
        att_level_rect = AttLevelRect.ElevDiff[np.minimum(elev_diff_matrix, 0) // -20, dist_matrix]
        # 修正初始结果
        att_level_results += att_level_rect
        att_level_results = np.maximum(att_level_results, 0).astype(int)  # 攻击等级不可能为负值
    # 根据攻击算子的位置概率，计算地图上所有六角格的加权攻击等级
    att_loc_probs_reduce = att_loc_probs[att_loc_probs > 0]
    threat_levels = att_loc_probs_reduce @ att_level_results
    return threat_levels


def calc_damage_common(myAgent, att_weapon_ids, att_bullet_nums, att_locs, tar_type, tar_locs,
                       att_obj_type=BopType.Vehicle, att_blood=3, armor_type=ArmorType.Composit, keep=0,
                       special_shoot=None):
    """ 2020-5-23 本函数较calc_damage.calc_damage_mean()更快速，支持矩阵计算
    用保存的期望损伤矩阵数据进行计算，支持输入参数为narray类型。主要用于预测损伤
    本函数中的attack_level参数为上方handle_shoot.calc_thread_att_level()函数的输出
    """
    if isinstance(att_locs, int):
        att_locs = np.array([att_locs])
    if isinstance(tar_locs, int):
        tar_locs = np.array([tar_locs])
    n_att_locs = len(att_locs)
    n_tar_locs = len(tar_locs)
    attack_level = calc_thread_att_level(myAgent, att_weapon_ids, att_bullet_nums, att_locs, tar_type, tar_locs,
                                         att_obj_type, att_blood, special_shoot)
    tar_terrains = myAgent.my_map['terrain_vec'][tar_locs]
    if tar_type == BopType.Vehicle:
        rect_values = np.zeros(n_tar_locs, dtype=int)
        for m in range(n_tar_locs):
            rect_values[m] = DamageRectVehicleDirect().calc_rect_value(tar_terrain=tar_terrains[m])  # 损伤结果修正只考虑目标地形
        damage_matrix = myAgent.damage_matrix[tar_type]
        return damage_matrix[attack_level, np.tile(armor_type, (n_att_locs, n_tar_locs)),
                             np.tile(att_blood - 1, (n_att_locs, n_tar_locs)),
                             np.tile(rect_values + 10, (n_att_locs, 1))]
    elif tar_type == BopType.Infantry:
        rect_values = np.zeros(n_tar_locs, dtype=int)
        for m in range(n_tar_locs):
            rect_values[m] = DamageRectInfantryDirect().calc_rect_value(tar_terrain=tar_terrains[m])  # 损伤结果修正只考虑目标地形
        damage_matrix = myAgent.damage_matrix[tar_type]
        return damage_matrix[attack_level, np.tile(keep, (n_att_locs, n_tar_locs)),
                             np.tile(rect_values + 10, (n_att_locs, 1))]
    elif tar_type == BopType.Aircraft:
        return DamageAircraft.KillProb[attack_level]  # 这里输出的是歼灭概率，上面2种情况是损伤车/班数的期望
    else:
        raise Exception('目标算子类型输入错误！')


def calc_damage_objs(myAgent, att_obj, tar_obj, special_shoot=None):
    """ 2020-5-23 本函数较calc_damage.calc_damage_mean()更快速，支持矩阵计算
    计算两两算子的期望损伤，主要用于实时计算损伤
    本函数中的attack_level参数为上方handle_shoot.calc_thread_att_level()函数的输出
    """
    att_locs = att_obj['cur_hex']
    tar_locs = tar_obj['cur_hex']
    att_weapon_ids = att_obj['carry_weapon_ids']
    att_bullet_nums = att_obj['remain_bullet_nums']
    tar_type = tar_locs['type']
    att_obj_type = att_obj['type']
    att_blood = att_obj['blood']
    armor_type = tar_obj['armor']
    tar_keep = tar_obj['keep']
    attack_level = calc_thread_att_level(myAgent, att_weapon_ids, att_bullet_nums, att_locs, tar_type, tar_locs,
                                         att_obj_type, att_blood, special_shoot)
    att_terrain = myAgent.my_map['terrain_vec'][att_locs]
    tar_terrain = myAgent.my_map['terrain_vec'][tar_locs]
    if tar_type == BopType.Vehicle:
        rect_value = DamageRectVehicleDirect(). \
            calc_rect_value(att_terrain, att_obj['keep'], att_obj['speed'], tar_terrain,
                            tar_obj['move_state'], tar_obj['speed'], tar_obj['stack'])
        damage_matrix = myAgent.damage_matrix[tar_type]
        return damage_matrix[attack_level - 1, armor_type, att_blood - 1, rect_value + 10]
    elif tar_type == BopType.Infantry:
        rect_value = DamageRectInfantryDirect(). \
            calc_rect_value(att_obj['keep'], att_obj['speed'], tar_terrain, tar_obj['move_state'],
                            tar_obj['speed'], tar_obj['stack'])
        damage_matrix = myAgent.damage_matrix[tar_type]
        return damage_matrix[attack_level - 1, tar_keep, rect_value + 10]
    elif tar_type == BopType.Aircraft:
        return DamageAircraft.KillProb[attack_level]  # 这里输出的是歼灭概率，上面2种情况是损伤车/班数的期望
    else:
        raise Exception('目标算子类型输入错误！')
